package com.example.ai.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import com.example.ai.config.tools.CurrentTimeTool;
import com.example.ai.config.tools.IpInfoTool;
import com.example.ai.dto.ChatReq;

import jakarta.annotation.Resource;
import reactor.core.publisher.Flux;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@RestController
class ChatController {
    @Resource
    private ChatClient chatClient;

    @PostMapping("/ai")
    String generation(@RequestBody ChatReq chatReq) {
        return this.chatClient.prompt()
                .advisors(a -> a
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatReq.chatId())
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .system(sp -> sp.param("voice", chatReq.msg()))
                .user(chatReq.msg())
                .tools(new CurrentTimeTool(), new IpInfoTool())
                .call()
                .content();
    }

    @PostMapping(value = "/ai-flux", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    Flux<String> generationStream(@RequestBody ChatReq chatReq) {
        return this.chatClient.prompt()
                .advisors(a -> a
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatReq.chatId())
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new SimpleLoggerAdvisor())
                .system(sp -> sp.param("voice", chatReq.msg()))
                .user(chatReq.msg())
                .tools(new CurrentTimeTool(), new IpInfoTool())
                .stream().content();
    }
}
